#coding:utf-8

import pandas as pd
import numpy as np
from stacking_model import load_excel_data, preprocess_data, train_stacking_model as train_stacking_func, predict_stacking_model as predict_stacking_func
import os

def prepare_and_train_model(train_file, target_col=None):
    """
    训练Stacking模型
    
    Args:
        train_file: 训练数据文件路径
        target_col: 目标列名，如果为None则使用最后一列
    
    Returns:
        训练好的模型状态字典
    """
    # 加载训练数据
    train_df = load_excel_data(train_file)
    
    # 使用统一的preprocess_data函数处理数据 - 传入任意非None值以确保返回特征和目标
    print("准备训练数据...")
    X_train, y_train = preprocess_data(train_df, target_col="train_mode")
    
    # 训练Stacking模型
    model_state = train_stacking_func(X_train, y_train)
    
    print("Stacking模型训练完成！")
    return model_state

def predict_with_stacking_model(model_state, predict_file, output_file):
    """
    使用训练好的Stacking模型进行预测
    
    Args:
        model_state: 训练好的模型状态字典
        predict_file: 预测数据文件路径
        output_file: 预测结果输出文件路径
    """
    # 加载预测数据
    predict_df = load_excel_data(predict_file)
    
    # 预处理预测数据 - 不指定target_col，只返回特征
    X_predict = preprocess_data(predict_df)
    
    # 进行预测 - 直接调用独立的预测函数
    print("开始预测...")
    predictions = predict_stacking_func(X_predict)
    
    # 保存预测结果
    result_df = predict_df.copy()
    result_df['预测结果'] = predictions
    
    try:
        result_df.to_excel(output_file, index=False, engine='openpyxl')
    except:
        try:
            result_df.to_excel(output_file, index=False, engine='xlrd')
        except Exception as e:
            print(f"保存预测结果失败: {e}")
            raise
    
    print(f"预测结果已保存到: {output_file}")
    print(f"预测结果形状: {result_df.shape}")
    return result_df

def main():
    """
    主函数，执行训练和预测流程
    """
    print("=== Stacking模型训练与预测程序 ===")
    
    # 文件路径设置
    train_file = '1.xlsx'  # 训练数据
    predict_file = '2.xlsx'  # 预测数据
    output_file = '3.xlsx'  # 输出结果
    
    # 检查文件是否存在
    for file in [train_file, predict_file]:
        if not os.path.exists(file):
            print(f"错误: 文件 '{file}' 不存在！")
            return
    
    try:
        # 训练模型
        model = prepare_and_train_model(train_file)
        
        # 进行预测
        result_df = predict_with_stacking_model(model, predict_file, output_file)
        
        print("\n=== 程序执行成功 ===")
        print(f"训练数据: {train_file}")
        print(f"预测数据: {predict_file}")
        print(f"输出结果: {output_file}")
        print("\n预测结果预览:")
        print(result_df.head())
        
    except Exception as e:
        print(f"程序执行出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()